package com.codingvip.tree.avl;

import com.sun.istack.internal.localization.NullLocalizable;

/**
 * @Author Tony.Han
 * @Created at 2020/4/6 9:01 上午
 * Desc:
 */
public class AVLTree {


    /**
     * @param root
     * @param data
     */
    public void insert(AVLNode root, int data) {
        if (root == null) {
            root = new AVLNode(data);
            return;
        }

        if (data < root.data) {
            if (root.left != null) {
                insert(root.left, data);
            } else {
                AVLNode newNode = new AVLNode(data);
                root.left = newNode;
                newNode.parent = root;

            }
        } else {
            if (root.right != null) {
                insert(root.right, data);
            } else {
                AVLNode newNode = new AVLNode(data);
                root.right = newNode;
                newNode.parent = root;
            }
        }
        root.balance = calBalance(root);
        boolean change = false;
        if (root.balance >= 2) {
            if (root.left.balance == -1) {
                leftRotate(root.left);
                change = true;
            }
            rightRotate(root);
            change = true;
        }

        if (root.balance <= -2) {
            if (root.left.balance == 1) {
                rightRotate(root.right);
                change = true;
            }
            leftRotate(root);
            change = true;
        }
        if (change) {
            root.balance = calBalance(root);
        }
        root.depth = calcDepth(root);


    }


    private void rightRotate(AVLNode current) {
        AVLNode pParent = current.parent;
        AVLNode pLeftSon = current.left;
        AVLNode pRightGrandSon = pLeftSon.right;

        //左子变父
        if (pParent != null) {
            if (current == pParent.left) {
                pParent.left = pLeftSon;
            } else if (current == pParent.right) {
                pParent.right = pLeftSon;
            }
        }
        pLeftSon.right = current;
        current.parent = pLeftSon;
        //右孙变左孙
        current.left = pRightGrandSon;
        if (pRightGrandSon != null) {
            pRightGrandSon.parent = current;
        }
        current.depth = calcDepth(current);
        current.balance = calBalance(current);
        pLeftSon.depth = calcDepth(pLeftSon);
        pLeftSon.balance = calBalance(pLeftSon);
    }

    private void leftRotate(AVLNode current) {
        AVLNode pParent = current.parent;
        AVLNode pRightSon = current.right;
        AVLNode pLeftGrandSon = pRightSon.left;

        //右子变父
        if (pParent != null) {
            if (current == pParent.left) {
                pParent.left = pRightSon;
            } else if (current == pParent.right) {
                pParent.right = pRightSon;
            }
        }
        pRightSon.left = current;
        current.parent = pRightSon;
        //左孙变右孙
        current.right = pLeftGrandSon;
        if (pLeftGrandSon != null) {
            pLeftGrandSon.parent = current;
        }
        current.depth = calcDepth(current);
        current.balance = calBalance(current);
        pRightSon.depth = calcDepth(pRightSon);
        pRightSon.balance = calBalance(pRightSon);

    }

    /**
     * @param current
     * @return
     */
    private int calcDepth(AVLNode current) {
        int depth = 0;
        if (current.left != null) {
            depth = current.left.depth;
        }
        if (current.right != null && depth < current.right.depth) {

            depth = current.right.depth;
        }
        depth++;
        return depth;
    }

    /**
     * @param current
     * @return
     */
    private int calBalance(AVLNode current) {
        int leftDepth;
        leftDepth = getDepthOfNode(current.left);
        int rightDepth = getDepthOfNode(current.right);
        return leftDepth - rightDepth;
    }

    private int getDepthOfNode(AVLNode root) {
        int leftDepth;
        if (root != null) {
            leftDepth = root.depth;
        } else {
            leftDepth = 0;
        }
        return leftDepth;
    }
}
